-- | Multivariate Bernstein Polynomials

module Bs (C(..), T, lo, hi, lohi, sign, ublb, norm) where

import qualified Tree
import qualified Ply
import qualified Iv
import Interlude
import qualified Data.Foldable as Foldable

-- | The class of multivariate Bernstein polynomials @b@ with
-- coefficients in @c@.
class C b where
  -- | Checks whether the polynomial is sharp w.r.t the given dimension.
  sharp :: Ord a => Int -> b a -> Maybe Hand
  -- | Projects the polynomial in the given dimension.
  proj :: Int -> Hand -> b a -> b a
  -- | Gives the derivative w.r.t the given dimension.
  diff :: Num a => Int -> b a -> b a
  -- | Subdivides into two parts
  subdiv :: Fractional a => Int -> b a -> (b a, b a)
  -- | Converts from the Horner basis.
  ofPly :: Fractional a => Ply.T a -> b a
  -- | Gives a tree representation. Used for metrcis.
  toTree :: Num a => b a -> Tree.T a
  -- | Evaluates the polynomial.
  eval :: Num a => b a -> [a] -> a

-- | The type of (unscaled) Bernstein polynomials.
newtype Bs a = B { unB :: Tree.T a } deriving (Eq,Show,Functor)

instance (Num a, Arbitrary a) => Arbitrary (Bs a) where
  arbitrary = B <$> Tree.genNonJagged arbitrary
  shrink = map B . shrink . unB

-- | Partial strict order
lt :: Ord a => Tree.T a -> Tree.T a -> Bool
lt (Tree.Lf x) (Tree.Lf y) = x < y
lt (Tree.Br x) (Tree.Br y) = all id $ zipWith lt x y

-- | Differentiation and integration in the first variable
diff0 :: {-Add-}Num a => a -> [a] -> [a]
diff0 _  []      = []
diff0 b0 (b1:b') = (b1 - b0) !: diff0 b1 b'

integ0 :: {-Add-}Num t => t -> [t] -> [t]
integ0 c []      = [c]
integ0 c (b0:b') = c !: integ0 (c + b0) b'

-- | Differential recursion
recL :: {-Add-}Num c => t -> (c -> [c] -> t -> t) -> [c] -> t
recL z f = r where
  r []      = z
  r (b0:b') = let d = diff0 b0 b' in f b0 d (r d)

instance C Bs where
  sharp = \k (B b) -> if sharpT k L b then Just L else
                      if sharpT k R b then Just R else Nothing where
    sharpT 0     L (Tree.Br (b0:b')) = all (lt b0) b'
    sharpT 0     R (Tree.Br b)       = sharpT 0 L (Tree.Br $ reverse b)
    sharpT (k+1) h (Tree.Br b)       = all (sharpT k h) b

  proj = \k h (B b) -> B $ projT k h b where
    projT 0     h (Tree.Br b) = (if h == L then head else last) b
    projT (k+1) h (Tree.Br b) = Tree.Br $ map (projT k h) b

  diff = \k (B b) -> B $ diffT k b where
    diffT 0     (Tree.Br (b0:b')) = Tree.Br $! diff0 b0 b'
    diffT (k+1) (Tree.Br b)       = Tree.Br $! map' (diffT k) b

  subdiv = \k (B b) -> (B $ subdivT L k b, B $ subdivT R k b) where
    subdivT L 0     (Tree.Br b) = Tree.Br $! subdiv0 b
    subdivT R 0     (Tree.Br b) = Tree.Br $! (reverse . subdiv0 . reverse) b
    subdivT h (k+1) (Tree.Br b) = Tree.Br $! map' (subdivT h k) b
    subdiv0 = recL [] (\c b -> integ0 (fmap ((2 ^ length b)*) c))

  eval = cjauT . unB where
    cjauT :: Num a => Tree.T a -> [a] -> a
    cjauT (Tree.Lf c) _      = c
    cjauT (Tree.Br b) (x:xs) = cjau0 x $ map (flip cjauT xs) b

    cjau0 :: Num a => a -> [a] -> a
    cjau0 _ [b0]     = b0
    cjau0 x b@(_:b') = cjau0 x $ zipWith (\b1 b2 -> (1-x)*b1 + x*b2) b b'

  ofPly = B . ofPlyT . balance . Ply.toTree where
    ofPlyT (Tree.Lf c) = Tree.Lf c
    ofPlyT p           = Tree.Br $! map ofPlyT $ ofPly0 p where
    ofPly0 = Ply.dRecT [] $ \c r ->
      integ0 (fmap ((fromInteger $ fact $ length r) *) c) r
    balance t = fill (Tree.dim t) t
    fill [] (Tree.Lf x) = Tree.Lf x
    fill (d0:d) (Tree.Lf x) = Tree.Br $ fill d (Tree.Lf x) : replicate (d0-1) (fill d (Tree.Lf 0))
    fill (d0:d) (Tree.Br x) = Tree.Br $ map (fill d) x ++ replicate (d0 - length x) (Tree.Lf 0)

  toTree = unB


-- | A Bernstein polynomial wrapped up with a scale factor.
data T a = Mk !(Bs a) !a deriving (Eq,Show)

instance C T where
  sharp k (Mk b _) = sharp k b
  proj k h (Mk b s) = Mk (proj k h b) s
  diff k (Mk b s) = Mk (diff k b) s

  subdiv k = \(Mk b s) -> let (l,r) = subdiv k b in
                          let s' = s * (0.5 ^ scale (unB b)) in
                            (Mk l s', Mk r s') where
    scale = (+ (-1)) . Tree.width . church k (\(Tree.Br (b:_)) -> b)

  eval (Mk b s) x = eval b x * s

  ofPly = (\b -> Mk b (recip $ fromInteger $ scale (unB b))) . ofPly where
    scale (Tree.Br (b0:b')) = fact (length b') * scale b0
    -- The scale factor, assuming the argument is regular (non-jagged).
    scale _            = 1

  toTree (Mk b s) = fmap (s*) $ toTree b

-- | Upper and lower bounds over the unit cube.

lo :: (Num a, Ord a) => T a -> a
lo (Mk (B b) s) = Tree.fold min b * s

hi :: (Num a, Ord a) => T a -> a
hi (Mk (B b) s) = Tree.fold max b * s

lohi :: (Num a, Ord a) => T a -> Iv.T a
lohi (Mk (B b) s) = Iv.mk (l*s) (h*s) where (l,h) = minmax b

sign :: (Num a, Ord a) => T a -> Maybe Bool
sign (Mk (B b) _) = sg $ Foldable.toList b where
  sg []     = Just True
  sg (x:xs) = case compare 0 x of
                EQ                -> sg xs
                LT | all (0<=) xs -> Just True
                GT | all (<=0) xs -> Just False
                _                 -> Nothing

corners :: Num a => T a -> [a]
corners (Mk (B b0) s) = cornersT b0 where
  cornersT (Tree.Lf c) = [c*s]
  cornersT (Tree.Br b) = cornersT (head b) `rapp` cornersT (last b)

-- | @ublb bs@ gives an upper bound on the lower bound of max
-- @bs@. Assumes that all polynomials in @bs@ have the same
-- dimensions.

ublb :: (Num a, Ord a) => [T a] -> a
ublb = foldl1' min . foldl1' (zipWith max) . map' corners

norm :: Num a => T a -> T a
norm (Mk b s) = Mk ((s*) <$> b) 1

{-# SPECIALIZE instance C T #-}
{-# SPECIALIZE ublb :: [T Double] -> Double #-}
